/**
 * 
 */
package edu.umd.clip.lm.storage.remote;

import java.io.*;
import java.net.*;
import java.nio.*;
import java.nio.channels.*;
import java.nio.channels.spi.SelectorProvider;
import java.text.DecimalFormat;
import java.text.NumberFormat;
import java.util.*;
import java.util.Map.Entry;
import java.util.concurrent.*;
import java.util.concurrent.atomic.*;
import java.util.concurrent.locks.*;

import edu.umd.clip.lm.util.*;

/**
 * @author Denis Filimonov <den@cs.umd.edu>
 *
 */
public class ServerStorageConnection implements Runnable {
	private SocketAddress address;
	private ServerSocketChannel channel;
	//private ConcurrentQueue<Task> tasks;
	private Selector selector;
	
	private ConcurrentHashMap<SocketChannel, ClientConnection> connections;
	//private ReadWriteLock connectionLock;
	
	private AtomicLong bytesSent;
	private AtomicLong bytesReceived;
	private AtomicLong messagesSent;
	private AtomicLong messagesReceived;
	private AtomicLong sendingTime;
	
	//private Set<Integer> outstandingMessages;
	private Lock lock;
	
	public ServerStorageConnection(SocketAddress address) throws IOException {
		this.address = address;
		channel = ServerSocketChannel.open();
		channel.configureBlocking(false);
		channel.socket().setReuseAddress(true);
		channel.socket().bind(address);
		
		selector = SelectorProvider.provider().openSelector();
		channel.register(selector, SelectionKey.OP_ACCEPT);
		
		lock = new ReentrantLock();
		
		//tasks = new ConcurrentQueue<Task>();
		
		bytesSent = new AtomicLong();
		bytesReceived = new AtomicLong();
		messagesSent = new AtomicLong();
		messagesReceived = new AtomicLong();
		sendingTime = new AtomicLong();
		
		connections = new ConcurrentHashMap<SocketChannel, ClientConnection>();
		
		//connectionLock = new ReentrantReadWriteLock();
		//outstandingMessages = Collections.synchronizedSet(new HashSet<Integer>(50));
	}
	
	public static class Task {
		private final SocketChannel socket;
		public final Message message;
		
		private Task(SocketChannel socket, Message message) {
			this.socket = socket;
			this.message = message;
		}
	}
	
	private class ClientConnection {
		private ConcurrentLinkedQueue<ByteBuffer> outputQueue;
		private SocketChannel channel;
		private Lock readLock;
		private Lock writeLock;
		
		private ClientConnection(SocketChannel socket) {
			channel = socket;
			outputQueue = new ConcurrentLinkedQueue<ByteBuffer>();
			readLock = new ReentrantLock();
			writeLock = new ReentrantLock();
		}
		
		// returns true if there's more data left to send
		private boolean send() throws IOException {
			while(true) {
				ByteBuffer buffer = outputQueue.peek();
				if (buffer == null) return false;
				
				int written = channel.write(buffer);
				if (written > 0) {
					messagesSent.incrementAndGet();
					bytesSent.addAndGet(written);
				}
				if (buffer.remaining() > 0) return true;
				outputQueue.poll();
			}
		}
		
		private void queueMessage(Message message) {
			ByteBuffer buffer = message.write();
			outputQueue.add(buffer);
		}
	}
	
	private Message readMessage(SocketChannel channel) throws IOException {
		try {
			Message message;
			message = Message.getMessage(channel);
			if (message != null) {
				//System.out.printf("Received message %d\n", message.getId());
				//System.err.println("Server Receiver: received a message");
				//outstandingMessages.add(message.getId());
				messagesReceived.incrementAndGet();
				bytesReceived.addAndGet(message.getSize());
			}
			return message;
			
		} catch(IOException e) {
			killConnection(channel);
			throw e;
		}
	}
	
	private Task makeTask(SocketChannel channel, Message message) throws IOException {
		if (message != null) {
			if (message.getType() == Message.MessageHeader.PING_MESSAGE) {
				Message pong = new PingMessage(message.getId(), 3000);
				ClientConnection connection = connections.get(channel);
				postReply(connection, pong);
			} else {
				Task task = new Task(channel, message);
				return task;
			}
		}
		return null;
	}

	public Task getNextTask() {
		while(true) {
			Message message = null;
			SocketChannel channel = null;

			//synchronized(selector) {
				lock.lock();
				try {
					selector.select();
				}catch(IOException e) {
					e.printStackTrace();
					return null;
				}
	
			
				Iterator<SelectionKey> iterator = selector.selectedKeys().iterator();

				while(true) {
					if (!iterator.hasNext()) {
						//lock.unlock();
						break;
					}

					SelectionKey key = iterator.next();
					iterator.remove();

					if (!key.isValid()) {
						continue;
					}
					//lock.unlock();

					if (key.isAcceptable()) {
						try {
							System.err.println("selected an acceptable key");
							acceptConnection((ServerSocketChannel) key.channel());
						} catch(IOException e) {
							e.printStackTrace();
						}
					} else if (key.isReadable()) {
						channel = (SocketChannel) key.channel();
						if (!channel.isOpen()) {
							key.cancel();
							break;
						}
						ClientConnection connection = connections.get(channel);
						connection.readLock.lock();
						// finished the message while the selector is blocked
						// otherwise another thread may attempt to finished this channel
						try {
							message = readMessage(channel);
						} catch(IOException e) {
							try {
								channel.close();
							} catch (IOException e1) {
							}
							key.cancel();
						} finally {
							connection.readLock.unlock();
						}
						break;
					} else if (key.isWritable()) {
						channel = (SocketChannel) key.channel();
						boolean hasMore = false;
						ClientConnection connection = connections.get(channel);
						try {
							connection.writeLock.lock();
							hasMore = connection.send();
						} catch(IOException e) {
							try {
								channel.close();
							} catch (IOException e1) {
							}
							key.cancel();						
						} finally {
							connection.writeLock.unlock();
						}

						if (!hasMore) {
							//lock.lock();
							key.interestOps(SelectionKey.OP_READ);
							selector.wakeup();
							//lock.unlock();
						}
					}
					break;
				}
				lock.unlock();
			//}

			if (message != null) {
				Task task = null;
				try {
					task = makeTask(channel, message);
				} catch (IOException e) {
					e.printStackTrace();
				}
				return task;
			}
		}
	}
	
	private void postReply(ClientConnection connection, Message message) {
		connection.queueMessage(message);
		boolean hasMore = false;
		try {
			connection.writeLock.lock();
			hasMore = connection.send();
		} catch (IOException e) {
			return;
		} finally {
			connection.writeLock.unlock();
		}
		
		if (hasMore) {
			lock.lock();
			SelectionKey key = connection.channel.keyFor(selector);
			key.interestOps(SelectionKey.OP_WRITE);

			selector.wakeup();
			lock.unlock();
		}		
	}
	
	public void postReply(Task task, Message message) {
		ClientConnection connection = connections.get(task.socket);
		if (connection != null) {
			postReply(connection, message);
		} else {
			System.err.println("can't find the connection to send a response to");
		}
	}
	
	private void killConnection(SocketChannel socket) {
		//connectionLock.writeLock().lock();
		ClientConnection connection = connections.remove(socket);
		if (connection != null) {
			lock.lock();
			SelectionKey key = socket.keyFor(selector);
			if (key != null) {
				selector.wakeup();
				key.cancel();
			}
			lock.unlock();
			//connection.stop();
		} else {
			System.err.printf("killing non-existent connection\n");
		}
		//connectionLock.writeLock().unlock();
	}
	
	private void acceptConnection(ServerSocketChannel channel) throws IOException {
		lock.lock();
		// Accept the connection and make it non-blocking
		SocketChannel socketChannel = channel.accept();
		if (socketChannel != null) {
			Socket socket = socketChannel.socket();
			socket.setTcpNoDelay(true);
			socket.setSendBufferSize(256*1024);
			socket.setReceiveBufferSize(64*1024);
			socketChannel.configureBlocking(false);
			
			// we won't be able to register while someone is doing select()
			selector.wakeup();
			socketChannel.register(selector, SelectionKey.OP_READ);

			System.err.println("Server: accepted a connection");
			ClientConnection connection = new ClientConnection(socketChannel);
			//connectionLock.writeLock().lock();
			connections.put(socketChannel, connection);
		}
		lock.unlock();
	}
	
	public void serverLoop() {
		while(true) {
			try {
				SocketChannel sock = channel.accept();
				sock.socket().setTcpNoDelay(true);
				sock.socket().setSendBufferSize(256*1024);
				sock.socket().setReceiveBufferSize(64*1024);
				System.err.println("Server: accepted a connection");
				ClientConnection connection = new ClientConnection(sock);
				//connectionLock.writeLock().lock();
				connections.put(sock, connection);
				//connection.start();
				//connectionLock.writeLock().unlock();
			} catch (IOException e) {
				// TODO Auto-generated catch block
				e.printStackTrace();
			}
		}
	}

	/* (non-Javadoc)
	 * @see java.lang.Runnable#run()
	 */
	@Override
	public void run() {
		serverLoop();
	}

	public long getBytesSent() {
		return bytesSent.longValue();
	}

	public long getBytesReceived() {
		return bytesReceived.longValue();
	}

	public long getMessagesSent() {
		return messagesSent.longValue();
	}

	public long getMessagesReceived() {
		return messagesReceived.longValue();
	}
	
	public long getSendingTime() {
		return sendingTime.longValue();
	}

	public static class ServerMonitor extends TimerTask {
		final ServerStorageConnection connection;
		long lastBytesSent = 0;
		long lastMessagesSent = 0;
		long lastBytesReceived = 0;
		long lastMessagesReceived = 0;
		long lastTimeSpent = 0;
		long lastSendingTime = 0;
		private NumberFormat nf;
	
		public static final long INTERVAL = 60000; 
		
		public ServerMonitor(ServerStorageConnection connection) {
			this.connection = connection;
			this.nf = new DecimalFormat("#.##");
		}
	
	
		/* (non-Javadoc)
		 * @see java.util.TimerTask#run()
		 */
		@Override
		public void run() {
			long bytesSent = connection.getBytesSent();
			long messagesSent = connection.getMessagesSent();
			long bytesReceived = connection.getBytesReceived();
			long messagesReceived = connection.getMessagesReceived();
			long timeSpent = ServerStorageResponder.getTimeSpend();
			long sendingTime = connection.getSendingTime();
			
			if (messagesSent != lastMessagesSent || messagesReceived != lastMessagesReceived) {
				System.err.printf("Server stats (1 min): IN %d KB/s %s req/s, OUT IN %d KB/s %s req/s, CPU %sms/req, net write %sms/req\n", 
						(bytesReceived - lastBytesReceived) * 1000 / INTERVAL / 1024,
						nf.format((float) (messagesReceived - lastMessagesReceived) / INTERVAL * 1000),
						(bytesSent - lastBytesSent) * 1000 / INTERVAL / 1024,
						nf.format((float) (messagesSent - lastMessagesSent) / INTERVAL * 1000),
						nf.format((timeSpent - lastTimeSpent) * 1e-6 / (messagesSent - lastMessagesSent)),
						nf.format((sendingTime - lastSendingTime) * 1e-6 / (messagesSent - lastMessagesSent)));
			}
			lastBytesSent = bytesSent;
			lastMessagesSent = messagesSent;
			lastBytesReceived = bytesReceived;
			lastMessagesReceived = messagesReceived;
			lastTimeSpent = timeSpent;
			lastSendingTime = sendingTime;
			
			//System.out.flush();
		} 
	}
}
